from abc import ABC,abstractclassmethod
import os
import re
from .base_extract import BaseExtract

"""
字符串匹配算法得到实体
"""
class MatchExtract(BaseExtract):
    """字典树匹配算法"""
    def __init__(self,entity_file_path):
        """
        Args:
            entity_file_path:实体列表
            格式如下
            实体1..类型1
            实体2..类型1
            ...
        """

        super(MatchExtract, self).__init__()

        self.trie_tree = {}
        self.entity2id = {}
        self.id2entity = {}
        self.type2id = {}
        self.entityname2type = {}
        self.end_sign = None

        self.ignore_set = {
            '体重',
            '精神',
            '手术',
            '不适',
            '休息',
            '明显增减',
            '出生时',
            'mg',
            '鸡蛋',
            '持续',
            '反复发作',
            '明显异常',
            '右侧',
            '第3',
            '全身',
            '明显变化',
            '受限',
            '转移',
            '刺激',
            '明显改变',
            '减轻',
            '大小',
            '量少',
            '中度',
            '为主',
            '突发',
            '呼吸',
            '明显减轻',
            '慢性',
            '量多',
            '脱出',
            '晚期',
            '渐增大',
            '重度',
            '中晚',
            '明显下降',
            '第三',
            '感染',
            '完全'
        }

        # 构建字典树
        self.build_tree(entity_file_path)
        self.entity_size = len(self.entity2id) # 实体数量
        self.distinguish_pos_neg = True


    def _add_entity_to_tree(self,entity_name:str):
        """
        添加实体到字典树中
        """
        tree_node = self.trie_tree
        # entity = re.split(r"[ ]",entity)
        for word in entity_name:
            if tree_node.get(word) is None:
                tree_node[word] = {}
            tree_node = tree_node[word]
        tree_node[self.end_sign] = None
    
    def _add_entity_to_list(self,entity_name:str,entity_type):
        """
        添加实体/类型到列表里面
        Args:
            entity_name:实体名字
            entity_type:实体类型
        """
        if entity_name not in self.entity2id:
            self.entity2id[entity_name] = len(self.entity2id)
            self.id2entity[len(self.id2entity)] = entity_name
        if entity_type not in self.type2id:
            self.type2id[entity_type] = len(self.type2id)
        self.entityname2type[entity_name] = entity_type
    
    def process_entity_name(self,entity_name):
        """
        标准化实体名字
        """
        entity_name = entity_name.lower()
        entity_name = entity_name.strip('等伴')
        return entity_name
    
    def _valid_entity_name(self, entity_name):
        """
        判断是否是合理的实体
        """
        if entity_name in self.ignore_set: return False
        if '无' in self.ignore_set: return False
        return True

    def build_tree(self,file_path):
        """
        建立字典树
        """
        with open(file_path,'r',encoding='utf-8') as f:
            for line in f:
                line = line.strip()
                if line == '': continue
                entity = line.split('\t')[0]
                entity_type = entity.split('..')[1]
                entity_name = entity.split('..')[0]
                entity_name = self.process_entity_name(entity_name)

                if self._valid_entity_name(entity_name):
                    self._add_entity_to_tree(entity_name)
                    self._add_entity_to_list(entity_name,entity_type)
    
    def _match_entity(self,start_idx:int,sentence:str) -> list:
        """
        从start位置匹配实体
        Args:
            start_idx:实体匹配位置
            sentence:文档
        Return:
            end_idxs:实体结束位置+1列表
        """
        tree_node = self.trie_tree
        token_idx = start_idx
        end_idxs = []
        while token_idx < len(sentence) and (sentence[token_idx] in tree_node):
            tree_node = tree_node[sentence[token_idx]]
            if self.end_sign in tree_node:
                end_idxs.append(token_idx + 1)
            token_idx += 1
        
        return end_idxs
    
    def _find_entities(self,sentence:str) -> list:
        """
        获取实体列表
        Args:
            sentence:句子/文档
        Return:
            entities:匹配实体列表 
            [(实体名, 实体类型, 实体开始位置, 实体结束位置),...]
        """
        start_idx = -1
        end_idx = -1
        entities = set()
        sentence = sentence.lower()
        # sentence = re.split(r'[\. ]',sentence)
        for start_idx,token in enumerate(sentence):
            end_idxs = self._match_entity(start_idx,sentence)
            if len(end_idxs) > 0:
                entity = sentence[start_idx:end_idxs[-1]]
                if len(entity) > 1: entities.add(entity)

        del_entities = set()
        for entity1 in entities:
            for entity2 in entities:
                if entity1 != entity2 and entity1 in entity2:
                    del_entities.add(entity1)
        
        entities = entities - del_entities

        return entities
    